#include "Pass.h"
#include "ComSubExprEli.h"
#include <set>
#include <algorithm>

static void CheckExprinst(exprinst& a);

void ComSubExprEli::execute()
{
    module->set_print_name();
    /*you need to finish this function*/


    for(auto fun: module->get_functions())
    {
        if (fun->get_basic_blocks().empty()) continue;
        func_ = fun;
        delete_in_blocks();
        create_inst_list();
        create_e_gen();
        calc_in_out();
        delete_global();
    }
}

bool ComSubExprEli::is_valid_expr(Instruction *inst)
{
    return !(inst->is_void()||inst->is_call()||inst->is_phi()||inst->is_alloca()||inst->is_load()||inst->is_cmp()||inst->is_zext());//TODO:CHECK VALID INST
}

void ComSubExprEli::create_inst_list()
{
    inst_list.clear();
    int len = 0;
    for (auto block : func_->get_basic_blocks())
        for (auto inst : block->get_instructions())
            if (is_valid_expr(inst)) 
            {
                if (dynamic_cast<BinaryInst*>(inst) == nullptr) continue;
                exprinst myinst;
                myinst.left = dynamic_cast<BinaryInst*>(inst)->get_operand(0);
                myinst.right = dynamic_cast<BinaryInst*>(inst)->get_operand(1);
                myinst.id = inst->get_instr_type();

                inst_list.push_back(myinst);
            }
}

void ComSubExprEli::delete_in_blocks()
{
    for(auto bb: func_->get_basic_blocks())
    {
        std::map<exprinst, User*> forward_list;
        std::map<User*, User*> new_value;
        std::set<Instruction*> delete_list;
        for(auto inst: bb->get_instructions())
        {
            if (!is_valid_expr(inst)) continue;
            if (dynamic_cast<BinaryInst*>(inst) == nullptr) continue;

            exprinst myinst;
            myinst.left = dynamic_cast<BinaryInst*>(inst)->get_operand(0);
            myinst.right = dynamic_cast<BinaryInst*>(inst)->get_operand(1);
            myinst.id = inst->get_instr_type();
            CheckExprinst(myinst);

            if (forward_list.find(myinst) == forward_list.end())
            {
                forward_list[myinst] = inst;
            }
            else
            {
                new_value[inst] = forward_list[myinst];
                delete_list.insert(inst);
            }
        }

        for(auto submap: new_value)
        {
            User* Oval = submap.first; 
            User* Nval = submap.second;
            for(auto use: Oval->get_use_list())
            {
                Instruction * use_inst = dynamic_cast<Instruction *>(use.val_);
                use_inst->set_operand(use.arg_no_, Nval);
            }
        } 
        for(auto inst:delete_list)
        {
            bb->delete_instr(inst);
        }       
    }
}

void ComSubExprEli::create_e_gen()
{
    e_gen.clear();
    int len = inst_list.size();
    int cnt = 0;
    for (auto block : func_->get_basic_blocks())
    {
        myset new_set(len);
        e_gen[block] = new_set;
        for (auto inst : block->get_instructions())
        {
            if (is_valid_expr(inst)) 
            {
                if (dynamic_cast<BinaryInst*>(inst) == nullptr) continue;
                
                exprinst myinst;
                myinst.left = dynamic_cast<BinaryInst*>(inst)->get_operand(0);
                myinst.right = dynamic_cast<BinaryInst*>(inst)->get_operand(1);
                myinst.id = inst->get_instr_type();
                CheckExprinst(myinst);

                e_gen[block].join(cnt++);
            }
        }
    }
}

void ComSubExprEli::calc_in_out()
{
    IN.clear();
    OUT.clear();
    int len = inst_list.size();
    bool flag = true;
    myset zero(len), one(len, true);
    
    for (auto block : func_->get_basic_blocks())
    {
        IN[block] = zero;
        OUT[block] = one;           
    }
    if (len == 0) return;
    while (flag)
    {
        flag = false;
        for (auto block : func_->get_basic_blocks())
        {
            myset tmp = one;
            if (block == func_->get_entry_block()) tmp = zero;
            
            for (auto pre_block : block->get_pre_basic_blocks())
            {
                tmp = tmp & OUT[pre_block];
            }
            
            IN[block] = tmp;
            myset new_out = e_gen[block] | tmp;
            if (!(OUT[block] == new_out)) flag = true;
            OUT[block] = new_out;
        }
    }
}

void ComSubExprEli::delete_global()
{
    std::map<User*, User*> new_value;
    std::map<BasicBlock*, std::set<Instruction*>> delete_list;
    for(auto bb: func_->get_basic_blocks())
    {
        for(auto inst: bb->get_instructions())
        {
            if (!is_valid_expr(inst)) continue;
            if (dynamic_cast<BinaryInst*>(inst) == nullptr) continue;

            exprinst myinst;
            myinst.left = dynamic_cast<BinaryInst*>(inst)->get_operand(0);
            myinst.right = dynamic_cast<BinaryInst*>(inst)->get_operand(1);
            myinst.id = inst->get_instr_type();
            CheckExprinst(myinst);

            int len = 0;
            Instruction* mem_inst = nullptr;
            for (auto bb2: func_->get_basic_blocks())
                for (auto inst2: bb2->get_instructions())
                {
                    if (!is_valid_expr(inst2)) continue;
                    if (dynamic_cast<BinaryInst*>(inst2) == nullptr) continue;
                    if (IN[bb].exist(len))
                    {
                        if (myinst == inst_list[len] && inst2 != inst)
                        {
                            mem_inst = inst2;
                            goto Break;
                        }
                    }
                    len++;
                }
Break:
            if (len != inst_list.size())
            {
                User* Nval = mem_inst;
                while (new_value.find(Nval) != new_value.end())
                    Nval = new_value[Nval];
                new_value[inst] = Nval;
                delete_list[bb].insert(inst);
            }
        } 
    }

    for(auto submap: new_value)
    {
        User* Oval = submap.first; 
        User* Nval = submap.second;
        while (new_value.find(Nval) != new_value.end())
            Nval = new_value[Nval];
        for(auto use: Oval->get_use_list())
        {
            Instruction * use_inst = dynamic_cast<Instruction *>(use.val_);
            use_inst->set_operand(use.arg_no_, Nval);
        }
    } 
    for(auto bb : func_->get_basic_blocks())
        for(auto inst : delete_list[bb])
        {
            bb->delete_instr(inst);
        }      
}

bool operator < (exprinst a, exprinst b)
{
    return 
        a.left < b.left ||
        a.left == b.left && a.right < b.right ||
        a.left == b.left && a.right == b.right && a.id < b.id;
}

bool operator == (exprinst a, exprinst b)
{
    return a.left == b.left && a.right == b.right && a.id == b.id;
}

void myset::join(int x)
{
    if (x > __size) return;
    __set[x / 32] |= 1 << x % 32;
}

bool myset::exist(int x)
{
    if (x > __size) return 0;
    return __set[x / 32] & (1 << x % 32);
}

myset operator & (myset& a, myset& b)
{
    int size = a.__size < b.__size ? a.__size : b.__size;
    myset c;
    c.__size = size;
    for (int i = 0; i < (size - 1) / 32 + 1; i++)
        c.__set.push_back(a.__set[i] & b.__set[i]);
    return c;
}

myset operator | (myset& a, myset& b)
{
    int size = a.__size < b.__size ? a.__size : b.__size;
    myset c;
    c.__size = size;
    for (int i = 0; i < (size - 1) / 32 + 1; i++)
        c.__set.push_back(a.__set[i] | b.__set[i]);
    return c;
}

bool operator == (myset& a, myset& b)
{
    if (a.__size != b.__size) return false;
    for (int i = 0; i < (a.__size - 1) / 32 + 1; i++)
        if (a.__set[i] != b.__set[i]) return false;
    return true;
}

static void CheckExprinst(exprinst& a)
{
    if (a.id == 
        Instruction::OpID::add || 
        Instruction::OpID::fadd ||
        Instruction::OpID::mul ||
        Instruction::OpID::fmul)
    {
        if (a.left > a.right) std::swap(a.left, a.right);
    }
}
